package io.unitycatalog.server.service;

import com.linecorp.armeria.common.HttpResponse;
import com.linecorp.armeria.common.HttpStatus;
import com.linecorp.armeria.server.annotation.Delete;
import com.linecorp.armeria.server.annotation.ExceptionHandler;
import com.linecorp.armeria.server.annotation.Get;
import com.linecorp.armeria.server.annotation.Param;
import com.linecorp.armeria.server.annotation.Patch;
import com.linecorp.armeria.server.annotation.Post;
import io.unitycatalog.server.auth.UnityCatalogAuthorizer;
import io.unitycatalog.server.auth.annotation.AuthorizeExpression;
import io.unitycatalog.server.auth.annotation.AuthorizeKey;
import io.unitycatalog.server.auth.annotation.AuthorizeKeys;
import io.unitycatalog.server.auth.decorator.UnityAccessEvaluator;
import io.unitycatalog.server.exception.GlobalExceptionHandler;
import io.unitycatalog.server.model.CatalogInfo;
import io.unitycatalog.server.model.CreateModelVersion;
import io.unitycatalog.server.model.CreateRegisteredModel;
import io.unitycatalog.server.model.FinalizeModelVersion;
import io.unitycatalog.server.model.ListRegisteredModelsResponse;
import io.unitycatalog.server.model.ModelVersionInfo;
import io.unitycatalog.server.model.RegisteredModelInfo;
import io.unitycatalog.server.model.SchemaInfo;
import io.unitycatalog.server.model.UpdateModelVersion;
import io.unitycatalog.server.model.UpdateRegisteredModel;
import io.unitycatalog.server.persist.CatalogRepository;
import io.unitycatalog.server.persist.MetastoreRepository;
import io.unitycatalog.server.persist.ModelRepository;
import io.unitycatalog.server.persist.Repositories;
import io.unitycatalog.server.persist.SchemaRepository;

import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.UUID;

import lombok.SneakyThrows;

import static io.unitycatalog.server.model.SecurableType.CATALOG;
import static io.unitycatalog.server.model.SecurableType.METASTORE;
import static io.unitycatalog.server.model.SecurableType.REGISTERED_MODEL;
import static io.unitycatalog.server.model.SecurableType.SCHEMA;

@ExceptionHandler(GlobalExceptionHandler.class)
public class ModelService extends AuthorizedService {

  private final ModelRepository modelRepository;
  private final SchemaRepository schemaRepository;
  private final CatalogRepository catalogRepository;
  private final MetastoreRepository metastoreRepository;

  private final UnityAccessEvaluator evaluator;

  @SneakyThrows
  public ModelService(UnityCatalogAuthorizer authorizer, Repositories repositories) {
    super(authorizer, repositories.getUserRepository());
    this.evaluator = new UnityAccessEvaluator(authorizer);
    this.catalogRepository = repositories.getCatalogRepository();
    this.schemaRepository = repositories.getSchemaRepository();
    this.modelRepository = repositories.getModelRepository();
    this.metastoreRepository = repositories.getMetastoreRepository();
  }

  @Post("")
  @AuthorizeExpression("""
      (#authorizeAny(#principal, #catalog, OWNER, USE_CATALOG) &&
          #authorize(#principal, #schema, OWNER)) ||
      (#authorizeAny(#principal, #catalog, OWNER, USE_CATALOG) &&
          #authorizeAll(#principal, #schema, USE_SCHEMA, CREATE_MODEL)) ||
      (#authorizeAny(#principal, #catalog, OWNER, USE_CATALOG) &&
          #authorizeAll(#principal, #schema, USE_SCHEMA, CREATE_FUNCTION))
      """)
  public HttpResponse createRegisteredModel(
      @AuthorizeKeys({
        @AuthorizeKey(value = SCHEMA, key = "schema_name"),
        @AuthorizeKey(value = CATALOG, key = "catalog_name")
      })
      CreateRegisteredModel createRegisteredModel) {
    assert createRegisteredModel != null;
    RegisteredModelInfo createRegisteredModelResponse =
        modelRepository.createRegisteredModel(createRegisteredModel);

    String catalogName = createRegisteredModelResponse.getCatalogName();
    String schemaName = createRegisteredModelResponse.getSchemaName();
    SchemaInfo schemaInfo = schemaRepository.getSchema(catalogName + "." + schemaName);
    String modelId = createRegisteredModelResponse.getId();
    initializeHierarchicalAuthorization(modelId, schemaInfo.getSchemaId());

    return HttpResponse.ofJson(createRegisteredModelResponse);
  }

  @Get("")
  @AuthorizeExpression("#defer")
  public HttpResponse listRegisteredModels(
      @Param("catalog_name") Optional<String> catalogName,
      @Param("schema_name") Optional<String> schemaName,
      @Param("max_results") Optional<Integer> maxResults,
      @Param("page_token") Optional<String> pageToken) {

    ListRegisteredModelsResponse listRegisteredModelsResponse =
        modelRepository.listRegisteredModels(catalogName, schemaName, maxResults, pageToken);
    filterModels("""
        #authorize(#principal, #metastore, OWNER) ||
        #authorize(#principal, #catalog, OWNER) ||
        (#authorize(#principal, #catalog, USE_CATALOG) && #authorize(#principal, #schema, OWNER)) ||
        (#authorizeAny(#principal, #registered_model, OWNER, EXECUTE) &&
            #authorize(#principal, #schema, USE_SCHEMA) &&
            #authorize(#principal, #catalog, USE_CATALOG))
        """, listRegisteredModelsResponse.getRegisteredModels());

    return HttpResponse.ofJson(listRegisteredModelsResponse);
  }

  @Get("/{full_name}")
  @AuthorizeExpression("""
      #authorize(#principal, #metastore, OWNER) ||
      #authorize(#principal, #catalog, OWNER) ||
      (#authorize(#principal, #catalog, USE_CATALOG) && #authorize(#principal, #schema, OWNER)) ||
      (#authorizeAny(#principal, #registered_model, OWNER, EXECUTE) &&
          #authorize(#principal, #schema, USE_SCHEMA) &&
          #authorize(#principal, #catalog, USE_CATALOG))
      """)
  @AuthorizeKey(METASTORE)
  public HttpResponse getRegisteredModel(
      @Param("full_name") @AuthorizeKey(REGISTERED_MODEL) String fullNameArg) {
    assert fullNameArg != null;
    RegisteredModelInfo registeredModelInfo = modelRepository.getRegisteredModel(fullNameArg);
    return HttpResponse.ofJson(registeredModelInfo);
  }

  @Patch("/{full_name}")
  @AuthorizeExpression("""
      (#authorize(#principal, #registered_model, OWNER) &&
          #authorizeAny(#principal, #schema, OWNER, USE_SCHEMA) &&
          #authorizeAny(#principal, #catalog, OWNER, USE_CATALOG))
      """)
  @AuthorizeKey(METASTORE)
  public HttpResponse updateRegisteredModel(
      @Param("full_name") @AuthorizeKey(REGISTERED_MODEL) String fullName,
      UpdateRegisteredModel updateRegisteredModel) {
    assert updateRegisteredModel != null;
    RegisteredModelInfo updateRegisteredModelResponse =
        modelRepository.updateRegisteredModel(fullName, updateRegisteredModel);
    return HttpResponse.ofJson(updateRegisteredModelResponse);
  }

  @Delete("/{full_name}")
  @AuthorizeExpression("""
      #authorize(#principal, #metastore, OWNER) ||
      #authorize(#principal, #catalog, OWNER) ||
      (#authorize(#principal, #catalog, USE_CATALOG) && #authorize(#principal, #schema, OWNER)) ||
      (#authorize(#principal, #registered_model, OWNER) &&
          #authorize(#principal, #schema, USE_SCHEMA) &&
          #authorize(#principal, #catalog, USE_CATALOG))
      """)
  @AuthorizeKey(METASTORE)
  public HttpResponse deleteRegisteredModel(
      @Param("full_name") @AuthorizeKey(REGISTERED_MODEL) String fullName,
      @Param("force") Optional<Boolean> force) {
    RegisteredModelInfo registeredModelInfo = modelRepository.getRegisteredModel(fullName);
    modelRepository.deleteRegisteredModel(fullName, force.orElse(false));

    SchemaInfo schemaInfo =
        schemaRepository.getSchema(
            registeredModelInfo.getCatalogName() + "." + registeredModelInfo.getSchemaName());
    removeHierarchicalAuthorizations(registeredModelInfo.getId(), schemaInfo.getSchemaId());

    return HttpResponse.of(HttpStatus.OK);
  }

  @Post("/versions")
  @AuthorizeExpression("""
      (#authorize(#principal, #registered_model, OWNER) &&
          #authorizeAny(#principal, #schema, OWNER, USE_SCHEMA) &&
          #authorizeAny(#principal, #catalog, OWNER, USE_CATALOG))
      """)
  public HttpResponse createModelVersion(
      @AuthorizeKeys({
        @AuthorizeKey(value = CATALOG, key = "catalog_name"),
        @AuthorizeKey(value = SCHEMA, key = "schema_name"),
        @AuthorizeKey(value = REGISTERED_MODEL, key = "model_name")
      })
      CreateModelVersion createModelVersion) {
    assert createModelVersion != null;
    assert createModelVersion.getModelName() != null;
    assert createModelVersion.getCatalogName() != null;
    assert createModelVersion.getSchemaName() != null;
    assert createModelVersion.getSource() != null;
    ModelVersionInfo createModelVersionResponse =
        modelRepository.createModelVersion(createModelVersion);
    return HttpResponse.ofJson(createModelVersionResponse);
  }

  @Get("/{full_name}/versions")
  @AuthorizeExpression("""
      #authorize(#principal, #metastore, OWNER) ||
      #authorize(#principal, #catalog, OWNER) ||
      (#authorize(#principal, #catalog, USE_CATALOG) && #authorize(#principal, #schema, OWNER)) ||
      (#authorizeAny(#principal, #registered_model, OWNER, EXECUTE) &&
          #authorize(#principal, #schema, USE_SCHEMA) &&
          #authorize(#principal, #catalog, USE_CATALOG))
      """)
  @AuthorizeKey(METASTORE)
  public HttpResponse listModelVersions(
      @Param("full_name") @AuthorizeKey(REGISTERED_MODEL) String fullName,
      @Param("max_results") Optional<Integer> maxResults,
      @Param("page_token") Optional<String> pageToken) {
    return HttpResponse.ofJson(modelRepository.listModelVersions(fullName, maxResults, pageToken));
  }

  @Get("/{full_name}/versions/{version}")
  @AuthorizeExpression("""
      #authorize(#principal, #metastore, OWNER) ||
      #authorize(#principal, #catalog, OWNER) ||
      (#authorize(#principal, #catalog, USE_CATALOG) && #authorize(#principal, #schema, OWNER)) ||
      (#authorizeAny(#principal, #registered_model, OWNER, EXECUTE) &&
          #authorize(#principal, #schema, USE_SCHEMA) &&
          #authorize(#principal, #catalog, USE_CATALOG))
      """)
  @AuthorizeKey(METASTORE)
  public HttpResponse getModelVersion(
      @Param("full_name") @AuthorizeKey(REGISTERED_MODEL) String fullName,
      @Param("version") Long version) {
    assert fullName != null && version != null;
    ModelVersionInfo modelVersionInfo = modelRepository.getModelVersion(fullName, version);
    return HttpResponse.ofJson(modelVersionInfo);
  }

  @Patch("/{full_name}/versions/{version}")
  @AuthorizeExpression("""
      (#authorize(#principal, #registered_model, OWNER) &&
          #authorizeAny(#principal, #schema, OWNER, USE_SCHEMA) &&
          #authorizeAny(#principal, #catalog, OWNER, USE_CATALOG))
      """)
  @AuthorizeKey(METASTORE)
  public HttpResponse updateModelVersion(
      @Param("full_name") @AuthorizeKey(REGISTERED_MODEL) String fullName,
      @Param("version") Long version,
      UpdateModelVersion updateModelVersion) {
    assert updateModelVersion != null;
    ModelVersionInfo updateModelVersionResponse =
        modelRepository.updateModelVersion(fullName, version, updateModelVersion);
    return HttpResponse.ofJson(updateModelVersionResponse);
  }

  @Delete("/{full_name}/versions/{version}")
  @AuthorizeExpression("""
      #authorize(#principal, #metastore, OWNER) ||
      #authorize(#principal, #catalog, OWNER) ||
      (#authorize(#principal, #catalog, USE_CATALOG) && #authorize(#principal, #schema, OWNER)) ||
      (#authorize(#principal, #registered_model, OWNER) &&
          #authorize(#principal, #schema, USE_SCHEMA) &&
          #authorize(#principal, #catalog, USE_CATALOG))
      """)
  @AuthorizeKey(METASTORE)
  public HttpResponse deleteModelVersion(
      @Param("full_name") @AuthorizeKey(REGISTERED_MODEL) String fullName,
      @Param("version") Long version) {
    modelRepository.deleteModelVersion(fullName, version);
    return HttpResponse.of(HttpStatus.OK);
  }

  @Patch("/{full_name}/versions/{version}/finalize")
  @AuthorizeExpression("""
      (#authorize(#principal, #registered_model, OWNER) &&
          #authorizeAny(#principal, #schema, OWNER, USE_SCHEMA) &&
          #authorizeAny(#principal, #catalog, OWNER, USE_CATALOG))
      """)
  @AuthorizeKey(METASTORE)
  public HttpResponse finalizeModelVersion(
      @Param("full_name") @AuthorizeKey(REGISTERED_MODEL) String fullName,
      FinalizeModelVersion finalizeModelVersion) {
    assert finalizeModelVersion != null;
    ModelVersionInfo finalizeModelVersionResponse =
        modelRepository.finalizeModelVersion(finalizeModelVersion);
    return HttpResponse.ofJson(finalizeModelVersionResponse);
  }

  public void filterModels(String expression, List<RegisteredModelInfo> entries) {
    // TODO: would be nice to move this to filtering in the Decorator response
    UUID principalId = userRepository.findPrincipalId();

    evaluator.filter(
        principalId,
        expression,
        entries,
        ti -> {
          CatalogInfo catalogInfo = catalogRepository.getCatalog(ti.getCatalogName());
          SchemaInfo schemaInfo =
              schemaRepository.getSchema(ti.getCatalogName() + "." + ti.getSchemaName());
          return Map.of(
              METASTORE,
              metastoreRepository.getMetastoreId(),
              CATALOG,
              UUID.fromString(catalogInfo.getId()),
              SCHEMA,
              UUID.fromString(schemaInfo.getSchemaId()),
              REGISTERED_MODEL,
              UUID.fromString(ti.getId()));
        });
  }
}

